﻿using System;
using System.Collections;
using System.Collections.Generic;
using System.Diagnostics.Contracts;
using System.Linq;

namespace Boggle.Net.Model
{
    /// <summary>
    /// Basic implementation of a trie data structure. Note that comparisons are case-sensitive.
    /// </summary>
    public sealed class Trie : IEnumerable<string>
    {
        private readonly Node _root = new Node(Char.MinValue);

        public Trie()
        {
        }

        public Trie(IEnumerable<string> words)
        {
            Contract.Requires<ArgumentNullException>(words != null);
            foreach (var word in words) Add(word);
        }

        public IEnumerator<string> GetEnumerator()
        {
            return AllWords.GetEnumerator();
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return GetEnumerator();
        }

        public void Add(string word)
        {
            Contract.Requires<ArgumentException>(!String.IsNullOrEmpty(word));

            var lastNode = word.Aggregate(_root, (node, letter) => node.Add(letter));
            lastNode.IsLast = true;
        }

        public bool Contains(string word)
        {
            var node = LastNodeOf(word);
            return node != null && node.IsLast;
        }

        public bool ContainsPrefix(string prefix)
        {
            return LastNodeOf(prefix) != null;
        }

        private Node LastNodeOf(string word)
        {
            Contract.Requires<ArgumentNullException>(word != null);

            if (word.Length == 0) return null;

            var node = _root;

            foreach (var letter in word)
            {
                var next = node.Find(letter);
                if (next == null) return null;

                node = next;
            }

            return node;
        }

        private IEnumerable<string> AllWords
        {
            get { return _root.Children.SelectMany(node => node.AllWordsStartingWith(String.Empty)); }
        }

        #region Embedded Types

        private class Node
        {
            private readonly ICollection<Node> _children = new List<Node>();
            private readonly char _letter;

            public Node(char letter)
            {
                _letter = letter;
            }

            public bool IsLast
            {
                get; set;
            }

            public IEnumerable<Node> Children
            {
                get { return _children; }
            }

            public Node Find(char letter)
            {
                return _children.FirstOrDefault(node => node._letter == letter);
            }

            public Node Add(char letter)
            {
                var node = Find(letter);

                if (node == null)
                {
                    node = new Node(letter);
                    _children.Add(node);
                }

                return node;
            }

            public IEnumerable<string> AllWordsStartingWith(string prefix)
            {
                var word = prefix + _letter;

                if (IsLast) yield return word;

                foreach (var extendedWord in Children.SelectMany(node => node.AllWordsStartingWith(word)))
                {
                    yield return extendedWord;
                }
            }
        }

        #endregion
    }
}
